package com.mahiru.mahiruaiagent.rag;

import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.rag.Query;
import org.springframework.ai.rag.preretrieval.query.transformation.QueryTransformer;
import org.springframework.ai.rag.preretrieval.query.transformation.RewriteQueryTransformer;
import org.springframework.stereotype.Component;

/**
 * @Author Enkidu
 * @Date 2025/7/7 16:22
 */
@Component
public class QueryRewriter {
    private final QueryTransformer queryTransformer;

    public QueryRewriter(ChatModel dashScopeChatModel) {
        ChatClient.Builder builder = ChatClient.builder(dashScopeChatModel);
        this.queryTransformer = RewriteQueryTransformer
                .builder()
                .chatClientBuilder(builder)
                .build();
    }

    /**
     * 对查询进行重写
     *
     * @param prompt
     * @return
     */
    public String doQueryRewrite(String prompt) {
        Query query = new Query(prompt);
        return queryTransformer.transform(query).text();
    }
}
